#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
脑电状态分析器

本文件用于实时分析用户的脑电状态，基于NeuroSky MindWave设备的原始数据。
与eeg_data_logger接轨，通过分析各频段功率、注意力和冥想值来判断用户当前状态。

支持的状态类型：
- 清醒状态：Beta波活跃，注意力值较高，冥想值较低
- 闭眼放松：Alpha波增强，冥想值上升，Beta波下降
- 专注状态：Beta/Gamma波活跃，注意力值高，冥想值适中
- 困倦状态：Delta/Theta波增强，注意力和冥想值都较低

分析原理：
1. 基于各频段功率的相对比例进行状态判断
2. 结合注意力和冥想指数进行辅助判断
3. 使用滑动窗口平滑处理，避免瞬时波动
4. 状态转换需要满足一定的持续时间要求
"""

import time
import csv
import os
import sys
from datetime import datetime
from collections import deque

# 添加父目录到路径以导入NeuroPy3
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from NeuroPy3 import NeuroPy

class EEGStateAnalyzer:
    """脑电状态分析器类"""
    
    def __init__(self, port="COM5", baudrate=57600, log_file="eeg_state_analysis.csv"):
        """
        初始化脑电状态分析器
        
        Args:
            port (str): 串口号，默认COM5
            baudrate (int): 波特率，默认57600
            log_file (str): 状态分析日志文件名
        """
        self.port = port
        self.baudrate = baudrate
        self.log_file = log_file
        self.neuropy = None
        self.is_analyzing = False
        
        # 当前数据存储
        self.current_data = {
            'timestamp': None,
            'attention': 0,
            'meditation': 0,
            'rawValue': 0,
            'delta': 0,
            'theta': 0,
            'lowAlpha': 0,
            'highAlpha': 0,
            'lowBeta': 0,
            'highBeta': 0,
            'lowGamma': 0,
            'midGamma': 0,
            'poorSignal': 0,
            'blinkStrength': 0
        }
        
        # 状态分析相关
        self.current_state = "未知"
        self.previous_state = "未知"
        self.state_duration = 0
        self.state_change_threshold = 3  # 状态改变需要持续3秒
        
        # 滑动窗口用于数据平滑（保存最近5个数据点）
        self.window_size = 5
        self.data_window = deque(maxlen=self.window_size)
        
        # 初始化日志文件
        self._init_csv_file()
    
    def _init_csv_file(self):
        """初始化CSV状态分析日志文件"""
        if not os.path.exists(self.log_file):
            with open(self.log_file, 'w', newline='', encoding='utf-8') as f:
                writer = csv.writer(f)
                headers = ['timestamp', 'attention', 'meditation', 'rawValue',
                          'delta', 'theta', 'lowAlpha', 'highAlpha',
                          'lowBeta', 'highBeta', 'lowGamma', 'midGamma',
                          'poorSignal', 'blinkStrength', 'analyzed_state',
                          'alpha_ratio', 'beta_ratio', 'gamma_ratio', 'state_confidence']
                writer.writerow(headers)
    
    def _calculate_frequency_ratios(self, data):
        """
        计算各频段的功率比例
        
        Args:
            data (dict): 包含各频段功率的数据字典
            
        Returns:
            tuple: (alpha_ratio, beta_ratio, gamma_ratio, total_power)
        """
        # 计算总功率（排除delta和theta，因为它们在清醒状态下通常较低）
        alpha_power = data['lowAlpha'] + data['highAlpha']
        beta_power = data['lowBeta'] + data['highBeta']
        gamma_power = data['lowGamma'] + data['midGamma']
        
        # 计算活跃频段总功率
        total_active_power = alpha_power + beta_power + gamma_power
        
        if total_active_power == 0:
            return 0, 0, 0, 0
        
        # 计算各频段比例
        alpha_ratio = alpha_power / total_active_power
        beta_ratio = beta_power / total_active_power
        gamma_ratio = gamma_power / total_active_power
        
        return alpha_ratio, beta_ratio, gamma_ratio, total_active_power
    
    def _analyze_state(self, data):
        """
        基于脑电数据分析用户状态
        
        Args:
            data (dict): 脑电数据字典
            
        Returns:
            tuple: (state, confidence)
        """
        attention = data['attention']
        meditation = data['meditation']
        
        # 计算频段比例
        alpha_ratio, beta_ratio, gamma_ratio, total_power = self._calculate_frequency_ratios(data)
        
        # 信号质量检查
        if data['poorSignal'] > 50:
            return "信号质量差", 0.1
        
        # 状态判断逻辑
        state = "未知"
        confidence = 0.5
        
        # 1. 闭眼放松状态：Alpha波占主导，冥想值高
        if (alpha_ratio > 0.5 and meditation > 80 and 
            attention < 70 and beta_ratio < 0.35):
            # 计算闭眼放松的置信度
            relax_score = 0.4  # 基础分
            if meditation > 90:
                relax_score += 0.3
            elif meditation > 80:
                relax_score += 0.2
            if alpha_ratio > 0.7:
                relax_score += 0.3
            elif alpha_ratio > 0.6:
                relax_score += 0.2
            if beta_ratio < 0.25:
                relax_score += 0.1
            
            state = "闭眼放松"
            confidence = min(0.9, relax_score)
        
        # 2. 专注状态：Beta/Gamma波活跃，注意力高
        elif (attention > 50 and meditation < 50 and 
              (beta_ratio > 0.25 or gamma_ratio > 0.2)):
            # 计算专注状态的置信度
            focus_score = 0
            if attention > 50:
                focus_score += 0.3
            if attention > 70:
                focus_score += 0.2
            if meditation < 50:
                focus_score += 0.2
            if gamma_ratio > 0.2:
                focus_score += 0.3
            if beta_ratio > 0.25:
                focus_score += 0.2
            
            state = "专注状态"
            confidence = min(0.8, max(0.5, focus_score))
        
        # 3. 清醒状态：Beta波适中，注意力和冥想值都不极端
        elif (attention > 40 and attention < 75 and 
              meditation < 70 and beta_ratio > 0.25 and beta_ratio < 0.6):
            state = "清醒状态"
            confidence = 0.7
        
        # 4. 放松状态：Alpha波增强但不如闭眼状态明显
        elif (alpha_ratio > 0.3 and meditation > 50 and 
              beta_ratio < 0.4 and attention < 70):
            state = "放松状态"
            confidence = 0.6 + alpha_ratio * 0.2
        
        # 5. 困倦状态：低频波增强，注意力和冥想都较低
        elif (attention < 40 and meditation < 50 and 
              (data['delta'] > total_power * 0.3 or data['theta'] > total_power * 0.3)):
            state = "困倦状态"
            confidence = 0.6
        
        # 6. 默认为清醒状态
        else:
            state = "清醒状态"
            confidence = 0.5
        
        return state, confidence
    
    def _smooth_analysis(self):
        """
        使用滑动窗口对状态分析进行平滑处理
        
        Returns:
            tuple: (smoothed_state, confidence)
        """
        if len(self.data_window) < 3:
            return self._analyze_state(self.current_data)
        
        # 对窗口内的数据进行分析
        state_votes = {}
        total_confidence = 0
        
        for data in self.data_window:
            state, confidence = self._analyze_state(data)
            if state in state_votes:
                state_votes[state] += confidence
            else:
                state_votes[state] = confidence
            total_confidence += confidence
        
        # 选择得票最高的状态
        if state_votes:
            best_state = max(state_votes, key=state_votes.get)
            avg_confidence = state_votes[best_state] / len(self.data_window)
            return best_state, min(0.9, avg_confidence)
        
        return "未知", 0.1
    
    def _log_analysis_to_csv(self, state, alpha_ratio, beta_ratio, gamma_ratio, confidence):
        """将分析结果记录到CSV文件"""
        if self.is_analyzing:
            timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]
            
            with open(self.log_file, 'a', newline='', encoding='utf-8') as f:
                writer = csv.writer(f)
                row = [
                    timestamp,
                    self.current_data['attention'],
                    self.current_data['meditation'],
                    self.current_data['rawValue'],
                    self.current_data['delta'],
                    self.current_data['theta'],
                    self.current_data['lowAlpha'],
                    self.current_data['highAlpha'],
                    self.current_data['lowBeta'],
                    self.current_data['highBeta'],
                    self.current_data['lowGamma'],
                    self.current_data['midGamma'],
                    self.current_data['poorSignal'],
                    self.current_data['blinkStrength'],
                    state,
                    f"{alpha_ratio:.3f}",
                    f"{beta_ratio:.3f}",
                    f"{gamma_ratio:.3f}",
                    f"{confidence:.3f}"
                ]
                writer.writerow(row)
    
    def _update_state(self, new_state, confidence):
        """更新当前状态，处理状态转换逻辑"""
        current_time = time.time()
        
        if new_state != self.current_state:
            # 状态发生变化，检查是否需要更新
            if new_state != self.previous_state:
                # 新的状态变化
                self.previous_state = new_state
                self.state_duration = 1
            else:
                # 持续相同的新状态
                self.state_duration += 1
                
                # 如果新状态持续足够长时间，则更新当前状态
                if self.state_duration >= self.state_change_threshold:
                    print(f"状态变化: {self.current_state} -> {new_state} (置信度: {confidence:.2f})")
                    self.current_state = new_state
                    self.state_duration = 0
        else:
            # 状态没有变化，重置计数器
            self.state_duration = 0
            self.previous_state = new_state
    
    # 回调函数定义
    def attention_callback(self, value):
        """注意力数据回调函数"""
        self.current_data['attention'] = value
        self._process_data()
    
    def meditation_callback(self, value):
        """冥想数据回调函数"""
        self.current_data['meditation'] = value
    
    def raw_value_callback(self, value):
        """原始脑电信号回调函数"""
        self.current_data['rawValue'] = value
    
    def delta_callback(self, value):
        """Delta波回调函数"""
        self.current_data['delta'] = value
    
    def theta_callback(self, value):
        """Theta波回调函数"""
        self.current_data['theta'] = value
    
    def low_alpha_callback(self, value):
        """低Alpha波回调函数"""
        self.current_data['lowAlpha'] = value
    
    def high_alpha_callback(self, value):
        """高Alpha波回调函数"""
        self.current_data['highAlpha'] = value
    
    def low_beta_callback(self, value):
        """低Beta波回调函数"""
        self.current_data['lowBeta'] = value
    
    def high_beta_callback(self, value):
        """高Beta波回调函数"""
        self.current_data['highBeta'] = value
    
    def low_gamma_callback(self, value):
        """低Gamma波回调函数"""
        self.current_data['lowGamma'] = value
    
    def mid_gamma_callback(self, value):
        """中Gamma波回调函数"""
        self.current_data['midGamma'] = value
    
    def poor_signal_callback(self, value):
        """信号质量回调函数"""
        self.current_data['poorSignal'] = value
    
    def blink_strength_callback(self, value):
        """眨眼强度回调函数"""
        self.current_data['blinkStrength'] = value
    
    def _process_data(self):
        """处理数据并进行状态分析"""
        # 将当前数据添加到滑动窗口
        self.data_window.append(self.current_data.copy())
        
        # 进行状态分析
        state, confidence = self._smooth_analysis()
        
        # 计算频段比例用于日志记录
        alpha_ratio, beta_ratio, gamma_ratio, _ = self._calculate_frequency_ratios(self.current_data)
        
        # 更新状态
        self._update_state(state, confidence)
        
        # 记录到日志文件
        self._log_analysis_to_csv(state, alpha_ratio, beta_ratio, gamma_ratio, confidence)
        
        # 实时显示当前状态
        print(f"当前状态: {self.current_state} | 分析状态: {state} | "
              f"注意力: {self.current_data['attention']} | 冥想: {self.current_data['meditation']} | "
              f"Alpha: {alpha_ratio:.2f} | Beta: {beta_ratio:.2f} | Gamma: {gamma_ratio:.2f}")
    
    def start_analysis(self):
        """开始脑电状态分析"""
        try:
            # 初始化NeuroPy对象
            self.neuropy = NeuroPy(self.port, self.baudrate)
            
            # 设置所有回调函数
            self.neuropy.setCallBack("attention", self.attention_callback)
            self.neuropy.setCallBack("meditation", self.meditation_callback)
            self.neuropy.setCallBack("rawValue", self.raw_value_callback)
            self.neuropy.setCallBack("delta", self.delta_callback)
            self.neuropy.setCallBack("theta", self.theta_callback)
            self.neuropy.setCallBack("lowAlpha", self.low_alpha_callback)
            self.neuropy.setCallBack("highAlpha", self.high_alpha_callback)
            self.neuropy.setCallBack("lowBeta", self.low_beta_callback)
            self.neuropy.setCallBack("highBeta", self.high_beta_callback)
            self.neuropy.setCallBack("lowGamma", self.low_gamma_callback)
            self.neuropy.setCallBack("midGamma", self.mid_gamma_callback)
            self.neuropy.setCallBack("poorSignal", self.poor_signal_callback)
            self.neuropy.setCallBack("blinkStrength", self.blink_strength_callback)
            
            # 开始数据采集和分析
            self.is_analyzing = True
            self.neuropy.start()
            
            print(f"开始脑电状态分析，结果保存到: {self.log_file}")
            print(f"串口: {self.port}, 波特率: {self.baudrate}")
            print("支持的状态: 清醒状态、闭眼放松、专注状态、放松状态、困倦状态")
            print("按 Ctrl+C 停止分析")
            print("-" * 80)
            
            # 保持程序运行
            while self.is_analyzing:
                time.sleep(1)
                
        except KeyboardInterrupt:
            print("\n用户中断，停止分析...")
            self.stop_analysis()
        except Exception as e:
            print(f"错误: {e}")
            self.stop_analysis()
    
    def stop_analysis(self):
        """停止脑电状态分析"""
        self.is_analyzing = False
        if self.neuropy:
            self.neuropy.stop()
        print("脑电状态分析已停止")

def main():
    """主函数"""
    # 创建状态分析器实例
    analyzer = EEGStateAnalyzer(port="COM5", log_file="eeg_state_analysis.csv")
    
    # 开始状态分析
    analyzer.start_analysis()

if __name__ == "__main__":
    main()